import os
# 解决OMP冲突: 必须在导入numpy或torch之前设置
os.environ['KMP_DUPLICATE_LIB_OK']='True'

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from tqdm import tqdm

from model import SimpleCNN
from data_loader import get_data_loaders

# --- 配置参数 ---
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DATA_DIR = './data'
MODEL_PATH = './gtsrb_cnn.pth'
BATCH_SIZE = 64
NUM_CLASSES = 43

def evaluate_model():
    """
    加载训练好的模型，并在测试集上进行评估。
    """
    print(f"使用设备: {DEVICE}")

    # 1. 检查模型文件是否存在
    if not os.path.exists(MODEL_PATH):
        print(f"错误: 找不到模型文件 '{MODEL_PATH}'。")
        print("请先运行 'python src/train.py' 来训练并保存模型。")
        return

    # 2. 加载数据
    try:
        _, test_loader = get_data_loaders(DATA_DIR, BATCH_SIZE)
        print("测试数据加载成功。")
    except Exception as e:
        print(f"数据加载失败: {e}")
        return

    # 3. 加载模型
    model = SimpleCNN(num_classes=NUM_CLASSES).to(DEVICE)
    model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
    model.eval()
    print(f"模型 '{MODEL_PATH}' 加载成功。")

    # 4. 在测试集上评估
    all_labels = []
    all_predictions = []

    eval_loop = tqdm(test_loader, desc="正在评估")
    with torch.no_grad():
        for images, labels in eval_loop:
            images, labels = images.to(DEVICE), labels.to(DEVICE)
            
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            
            all_labels.extend(labels.cpu().numpy())
            all_predictions.extend(predicted.cpu().numpy())

    # 5. 计算并打印评估指标
    print("\n--- 模型评估结果 ---")

    # 总体准确率
    accuracy = accuracy_score(all_labels, all_predictions)
    print(f"\n总体准确率 (Overall Accuracy): {accuracy * 100:.2f}%\n")

    # 分类报告 (Precision, Recall, F1-score)
    print("分类报告:")
    # target_names 可以根据数据集的类别说明进行修改
    # target_names = [f'Class {i}' for i in range(NUM_CLASSES)]
    # print(classification_report(all_labels, all_predictions, target_names=target_names))
    print(classification_report(all_labels, all_predictions))


    # 混淆矩阵
    print("正在生成混淆矩阵...")
    cm = confusion_matrix(all_labels, all_predictions)
    plt.figure(figsize=(20, 20))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=range(NUM_CLASSES), yticklabels=range(NUM_CLASSES))
    plt.xlabel('预测类别 (Predicted Label)')
    plt.ylabel('真实类别 (True Label)')
    plt.title('混淆矩阵 (Confusion Matrix)')
    
    # 保存混淆矩阵图像
    cm_path = 'confusion_matrix.png'
    plt.savefig(cm_path)
    print(f"\n混淆矩阵已保存为 '{cm_path}'")
    plt.show()


if __name__ == '__main__':
    evaluate_model() 